#include <Windows.h>
#include <stdio.h>
#define JMP( frm, to ) ( ( ( int )to - ( int )frm ) - 5 )

HWND    hwndClient = NULL;
HMODULE hLibModule = NULL;
void*   SendRet    = NULL;
void*   RecvRet    = NULL;

void __stdcall SendPacket( DWORD nSize, PVOID lpData ) {
  COPYDATASTRUCT cds = {0};

  cds.cbData = nSize;
  cds.lpData = lpData;
  cds.dwData = 1;

  SendMessage( hwndClient, WM_COPYDATA, 0, ( LPARAM )&cds );
}

void __stdcall RecvPacket( DWORD nSize, PVOID lpData ) {
  COPYDATASTRUCT cds = {0};

  cds.cbData = nSize;
  cds.lpData = lpData;
  cds.dwData = 0;

  SendMessage( hwndClient, WM_COPYDATA, 0, ( LPARAM )&cds );
}

__declspec( naked ) void SendHookProc() {
  __asm {
    push ebp
    mov ebp, esp

    mov eax, [ebp+0x0C]
    mov ecx, [ebp+0x10]
    push eax
    push ecx

    call SendPacket

    jmp [SendRet]
  }
}

__declspec( naked ) void RecvHookProc() {
  __asm {
    push ebp
    mov ebp, esp

    mov eax, [ebp+0x0C]
    mov ecx, [ebp+0x10]
    push eax
    push ecx

    call RecvPacket

    jmp [RecvRet]
  }
}

void WriteJMP( LPVOID from, LPVOID to ) {
  DWORD flOldProtect = 0;

  VirtualProtect( from, 5, PAGE_EXECUTE_READWRITE, &flOldProtect );
  *( BYTE* )from                 = 0xE9;
  *( DWORD* )( (byte*)from + 1 ) = JMP( from, to );
  VirtualProtect( from, 5, flOldProtect, &flOldProtect );
}

void RestoreHookBytes( LPVOID lpFunc ) {
  BYTE  origBytes[]  = {0x8B, 0xFF, 0x55, 0x8B, 0xEC};
  DWORD flOldProtect = 0;

  VirtualProtect( lpFunc, sizeof(origBytes), PAGE_EXECUTE_READWRITE, &flOldProtect );
  memcpy_s( lpFunc, sizeof(origBytes), origBytes, sizeof(origBytes) );
  VirtualProtect( lpFunc, sizeof(origBytes), flOldProtect, &flOldProtect );
}

void Init( HWND hwndInit ) {
  HMODULE hWinsock = GetModuleHandle( L"Ws2_32.dll" );
  hwndClient       = hwndInit;

  if( hWinsock == NULL ) {
    MessageBox( NULL, L"Winsock library non-existent", NULL, MB_OK | MB_ICONEXCLAMATION );
    FreeLibraryAndExitThread( hLibModule, 0 );
  } else {
    void* lpSend;
    void* lpRecv;

    lpSend  = GetProcAddress( hWinsock, "send" );
    SendRet = (byte*)lpSend + 5;
    WriteJMP( lpSend, SendHookProc );

    lpRecv  = GetProcAddress( hWinsock, "recv" );
    RecvRet = (byte*)lpRecv + 5;
    WriteJMP( lpRecv, RecvHookProc );
  }
}

void Cleanup() {
  HMODULE hWinsock = GetModuleHandle( L"Ws2_32.dll" );

  if( hWinsock != NULL ) {
    RestoreHookBytes( GetProcAddress( hWinsock, "send" ) );
    RestoreHookBytes( GetProcAddress( hWinsock, "recv" ) );
  }

  FreeLibraryAndExitThread( hLibModule, 0 );
}

BOOL WINAPI DllMain( HANDLE hinstDLL, DWORD dwReason, LPVOID lpvReserved ) {
  switch( dwReason ) {
    case DLL_PROCESS_ATTACH: {
      hLibModule = ( HMODULE )hinstDLL;
      break;
    }
    default:
      return FALSE;
  }

  return TRUE;
}